#include <math.h>
#include <fstream>
#include <iostream>
#include <amp.h>
#include <amp_math.h>
#include <ppl.h>
#include "mtgp32-fast.h"

//#pragma warning (disable : 4267)

using namespace concurrency;

class MtgpSimple
{
public:

    ////////////////////////////////////////////////////////////////////////////////
    // Transform each of MT_RNG_COUNT lanes of n_per_RNG uniformly distributed 
    // random samples, produced by rand_MT_amp(), to normally distributed lanes
    // using Cartesian form of Box-Muller transformation.
    ////////////////////////////////////////////////////////////////////////////////
    static void box_muller_transform(float& u1, float& u2) restrict(cpu, amp)
    {
        float r = fast_math::sqrt(-2.0f * fast_math::log(u1));
        float phi = 2.0f * 3.14159265358979f * u2;
        u1 = r * fast_math::cos(phi);
        u2 = r * fast_math::sin(phi);
    }

    static void box_muller_kernel(
        int idx, 
        array_view<float, 2> random, 
        int n_per_RNG
    ) restrict(cpu, amp)
    {
        for(int out = 0; out < n_per_RNG; out += 2) 
        {
            box_muller_transform(
                random[index<2>(idx, out)], 
                random[index<2>(idx, out+1)]
            );
        }
    }

    void generate_rand_on_amp(
        array_view<float,2> random, 
        bool apply_transform)
    {
        int seed = 1;
        int extent1 = _randomsExtent[1];
        array_view<const mtgp32_params_fast_t> mtgp_params_array(
            sizeof(mtgp32_params_fast_44497)/sizeof(mtgp32_params_fast_t),
            mtgp32_params_fast_44497);

        mtgp32_status_fast_t st;
        int size, large_size;
        get_sizes(mtgp32_params_fast_44497, size, large_size);
        st.idx = size - 1;
        st.size = size;
        st.large_size = large_size;
        st.large_mask = large_size -1;
        vector<mtgp32_status_fast_t> h_status(_randomsExtent[0], st); //create n copies of status 
        array<mtgp32_status_fast_t> d_status(_randomsExtent[0], h_status.begin(), _acc); //create on device

        Concurrency::parallel_for_each(_acc, extent<1>(_randomsExtent[0]), [=, &d_status] (index<1> idx) restrict(amp)
        {
            const mtgp32_params_fast_t& params = mtgp_params_array[idx[0]];
            mtgp32_status_fast_t* status = &d_status[idx[0]];
            mtgp32_init_state(status->array, &params, seed);
            for (int i = 0; i < extent1; i++) 
            {
                //direct3d_printf("d %d i %d\n", idx[0], status->idx);
                random(idx[0], i) = mtgp32_genrand_open_open(params, status);
            }
        });

        if (apply_transform)
        {
            Concurrency::parallel_for_each(extent<1>(_randomsExtent[0]), [=] (index<1> idx) restrict(amp)
            {
                box_muller_kernel(idx[0], random, extent1);
            });
        }
    }

    void generate_rand_on_cpu(
        std::vector<float>& v_random, 
        bool apply_transform)
    {
        array_view<float, 2> random(_randomsExtent, v_random);
        int extent1 = _randomsExtent[1];
        int seed = 1;

        mtgp32_status_fast_t st;
        int size, large_size;
        get_sizes(mtgp32_params_fast_44497, size, large_size);
        st.idx = size - 1;
        st.size = size;
        st.large_size = large_size;
        st.large_mask = large_size -1;
        vector<mtgp32_status_fast_t> h_status(_randomsExtent[0], st); //create n copies of status 

        // generate random numbers
        Concurrency::parallel_for(0, _randomsExtent[0], [&] (int idx) restrict(cpu)
        {
            mtgp32_params_fast_t *params = mtgp32_params_fast_44497 + idx;
            mtgp32_status_fast_t* status = &h_status[idx];
            mtgp32_init_state(status->array, params, seed);
            //mtgp32_fast_t mtgp32;
            //int mtgp_ret = mtgp32_init(&mtgp32, params, seed);
            //if (mtgp_ret) {
            //    throw new std::exception("failure in mtgp32_init\n");
            //}
            for (int i = 0; i < _randomsExtent[1]; i++) 
            {
                //mtgp32_status_fast_t* status = mtgp32.status;
                //printf("h %d i %d\n", idx, status->idx);
                random(idx, i) = mtgp32_genrand_open_open(*params, status);
            }
            //mtgp32_free(&mtgp32);
        });

        if (apply_transform)
        {
            Concurrency::parallel_for(0, _randomsExtent[0], [&] (int idx) restrict(cpu)
            {
                box_muller_kernel(idx, random, _randomsExtent[1]);
            });
        }
    }

    MtgpSimple(const extent<2>& randomsExtent, accelerator_view acc)
        :
        _mexp(44497),  //or 23209 or 110503
        _randomsExtent(randomsExtent),
        _acc(acc)
    {
        int no = 5; //no must be between 0 and 127
        uint32_t seed = 1;
        //uint32_t seed_ar[4] = {1, 2, 3, 4};
        //char seed_str[] = "\01\02\03\04";

        // Copy to GPU
        //d_matrix_a = array<unsigned int, 1>(inputsExtent, h_matrix_a.begin());
    }

    int _mexp;
    extent<2> _randomsExtent;
    accelerator_view _acc;
};
